{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameter Management"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The ultimate goal of training deep networks is to find good parameter values for a given architecture. When everything is standard, the __torch.nn.Sequential__ class is a perfectly good tool for it. However, very few models are entirely standard and most scientists want to build things that are novel. This section shows how to manipulate parameters. In particular we will cover the following aspects:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- Accessing parameters for debugging, diagnostics, to visualize them or to save them is the first step to understanding how to work with custom models.\n",
    "- Secondly, we want to set them in specific ways, e.g. for initialization purposes. We discuss the structure of parameter initializers.\n",
    "- Lastly, we show how this knowledge can be put to good use by building networks that share some parameters."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As always, we start from our trusty Multilayer Perceptron with a hidden layer. This will serve as our choice for demonstrating the various features."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "net = nn.Sequential()\n",
    "net.add_module('Linear_1', nn.Linear(20, 256, bias = False))\n",
    "net.add_module('relu', nn.ReLU())\n",
    "net.add_module('Linear_2', nn.Linear(256, 10, bias = False))\n",
    "\n",
    "# the init_weights function initializes the weights of our multi-layer perceptron \n",
    "\n",
    "def init_weights(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "\n",
    "# the net.apply() applies the above stated initialization of weights to our net        \n",
    "\n",
    "net.apply(init_weights) \n",
    "\n",
    "x = torch.randn(2,20)   #initialing a random tensor of shape (2,20)\n",
    "net(x)  #Forward computation\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter Access\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the case of a Sequential class we can access the parameters with ease, simply by calling __net.parameters__. Let’s try this out in practice by inspecting the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(net[0].parameters)\n",
    "print(net[1].parameters)\n",
    "print(net[2].parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The output tells us a number of things. Firstly, there are 3 layers; 2 linear layers and 1 ReLU layer as we would expect. The output also specifies the shapes that we would expect from linear layers. In particular the names of the parameters are very useful since they allow us to identify parameters uniquely even in a network of hundreds of layers and with nontrivial structure. Also, the output tells us that bias is __False__ as we specified it.  \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Targeted Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to do something useful with the parameters we need to access them, though. There are several ways to do this, ranging from simple to general. Let’s look at some of them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(net[0].bias)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It returns the bias of the first linear layer. Since we initialized the bias to be __False__, the output is None. We can also access the parameters by name, such as **Linear_1**. Both methods are entirely equivalent but the first method leads to much more readable code."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(net.Linear_1.weight)\n",
    "print(net[0].weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the weights are nonzero. This is by design since we applied __Xavier initialization__ to our network. We can also compute the gradient with respect to the parameters. It has the same shape as the weight. However, since we did not invoke backpropagation yet, the output is None."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(net[0].weight.grad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### All parameters at once"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Accessing parameters as described above can be a bit tedious, in particular if we have more complex blocks, or blocks of blocks (or even blocks of blocks of blocks), since we need to walk through the entire tree in reverse order to how the blocks were constructed. To avoid this, blocks come with a method __state_dict__ which grabs all parameters of a network in one dictionary such that we can traverse it with ease. It does so by iterating over all constituents of a block and calls __state_dict__ on subblocks as needed. To see the difference consider the following:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(net[0].state_dict) # parameters only for first layer\n",
    "print(net.state_dict) # parameters for entire network"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This provides us with a third way of accessing the parameters of the network. If we wanted to get the value of the weight term of the second linear layer we could simply use this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net.state_dict()['Linear_1.weight']"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Rube Goldberg strikes again"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's see how the parameter naming conventions work if we nest multiple blocks inside each other. For that we first define a function that produces blocks (a block factory, so to speak) and then we combine these inside yet larger blocks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def block1():\n",
    "    net = nn.Sequential(nn.Linear(16, 32),\n",
    "                        nn.ReLU(),\n",
    "                        nn.Linear(32, 16),\n",
    "                        nn.ReLU())\n",
    "    return net\n",
    "\n",
    "def block2():\n",
    "    net = nn.Sequential()\n",
    "    for i in range(4):\n",
    "        net.add_module('block'+str(i), block1())\n",
    "    return net    \n",
    "        \n",
    "rgnet = nn.Sequential()\n",
    "rgnet.add_module('model',block2())\n",
    "rgnet.add_module('Last_linear_layer', nn.Linear(16,10))\n",
    "rgnet.apply(init_weights)\n",
    "x = torch.randn(2,16)\n",
    "rgnet(x) # forward computation\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we are done designing the network, let's see how it is organized. __state_dict__ provides us with this information, both in terms of naming and in terms of logical structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(rgnet.parameters)\n",
    "for param in rgnet.parameters():\n",
    "    print(param.size(), param.dtype) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the layers are hierarchically generated, we can also access them accordingly. For instance, to access the first major block, within it the second subblock and then within it, in turn the bias of the first layer, we perform the following."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rgnet[0][1][0].bias.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Parameter Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we know how to access the parameters, let's look at how to initialize them properly. We discussed the need for initialization in section Numerical Stability. We often need to use methods to initialize the weights. PyTorch's init module provides a variety of preset initialization methods, but if we want something out of the ordinary, we need a bit of extra work. To initialize the weights of a single layer, we use a function from __torch.nn.init__ . For instance:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "linear1 = nn.Linear(2,5,bias=True)\n",
    "torch.nn.init.normal_(linear1.weight, mean=0, std =0.01)  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we wanted to initialize all parameters to 1, we could do this simply by changing the initializer to `Constant(1)`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def init_weight(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.normal_(m.weight)\n",
    "        \n",
    "net = nn.Sequential()\n",
    "net.add_module('Linear_1', nn.Linear(2, 5, bias = False))\n",
    "net.add_module('Linear_2', nn.Linear(5, 5, bias = False))\n",
    "\n",
    "net.apply(init_weight)\n",
    "print(net.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Built-in Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let’s begin with the built-in initializers. The code below initializes all parameters with Gaussian random variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gaussian_normal(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.normal_(m.weight)\n",
    "        \n",
    "net.apply(gaussian_normal)\n",
    "print(net[0].weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we wanted to initialize all parameters to 1, we could do this simply by changing the initializer to __torch.nn.init.constant_(tensor,1)__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ones(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.constant_(m.weight, 1)\n",
    "        \n",
    "net.apply(ones)\n",
    "print(net.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to initialize only a specific parameter in a different manner, we can simply set the initializer only for the appropriate subblock (or parameter) for that matter. For instance, below we initialize the __second layer__ to a constant value of __42__ and we use the __Xavier initializer__ for the weights of the __first layer__."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "block1 = nn.Sequential()\n",
    "block1.add_module('Linear_1', nn.Linear(2,5,bias=False))\n",
    "block2 = nn.Sequential()\n",
    "block2.add_module('Linear_2', nn.Linear(5,5,bias=False))\n",
    "\n",
    "model = nn.Sequential()\n",
    "model.add_module('first', block1)\n",
    "model.add_module('second', block2)\n",
    "\n",
    "def xavier_normal(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.xavier_uniform_(m.weight)\n",
    "def init_42(m):\n",
    "    if type(m) == nn.Linear:\n",
    "        torch.nn.init.constant_(m.weight, 42)\n",
    "              \n",
    "block1.apply(xavier_normal)\n",
    "block2.apply(init_42)\n",
    "print(model.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Initialization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, the initialization methods we need are not provided in the init module. At this point, we can implement our desired implementation by writing the desired functions and use them to initialize the weights. In the example below, we pick a decidedly bizarre and nontrivial distribution, just to prove the point. We draw the coefficients from the following distribution:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "$$ \\begin{aligned} w \\sim \\begin{cases} U[5, 10] & \\text{ with probability } \\frac{1}{4} \\   \n",
    "0 & \\text{ with probability } \\frac{1}{2} \\      \n",
    "U[-10, -5] & \\text{ with probability } \\frac{1}{4} \\end{cases} \\end{aligned} $$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def custom(m):\n",
    "    torch.nn.init.uniform_(m[0].weight, -10,10)\n",
    "    for i in range(m[0].weight.data.shape[0]):\n",
    "        for j in range(m[0].weight.data.shape[1]):\n",
    "            if m[0].weight.data[i][j]<=5 and m[0].weight.data[i][j]>=-5:\n",
    "                m[0].weight.data[i][j]=0\n",
    "    \n",
    "    \n",
    "m = nn.Sequential(nn.Linear(5,5,bias=False))\n",
    "custom(m)\n",
    "print(m.state_dict())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If even this functionality is insufficient, we can set parameters directly. Since __.data__ returns a Tensor we can access it just like any other matrix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "m[0].weight.data +=1\n",
    "m[0].weight.data[0][0] = 42\n",
    "m[0].weight.data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tied Parameters"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In some cases, we want to share model parameters across multiple layers. For instance when we want to find good word embeddings we may decide to use the same parameters both for encoding and decoding of words. Let’s see how to do this a bit more elegantly. In the following we allocate a linear layer and then use it multiple times for sharing the weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We need to give the shared layer a name such that we can reference its\n",
    "# parameters\n",
    "\n",
    "shared = nn.Sequential()\n",
    "shared.add_module('linear_shared', nn.Linear(8,8,bias=False))\n",
    "shared.add_module('relu_shared', nn.ReLU())                  \n",
    "net = nn.Sequential(nn.Linear(20,8,bias=False),\n",
    "               nn.ReLU(),\n",
    "               shared,\n",
    "               shared,\n",
    "               nn.Linear(8,10,bias=False))\n",
    "\n",
    "net.apply(init_weights)\n",
    "\n",
    "print(net[2][0].weight==net[3][0].weight)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above example shows that the parameters of the second and third layer are tied. They are identical rather than just being equal. That is, by changing one of the parameters the other one changes, too. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "* We have several ways to access, initialize, and tie model parameters.\n",
    "* We can use custom initialization.\n",
    "* PyTorch has a sophisticated mechanism for accessing parameters in a unique and hierarchical manner."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercises\n",
    "\n",
    "1. Use the FancyMLP defined in :numref:`chapter_model_construction` and access the parameters of the various layers.\n",
    "1. Look at the [PyTorch documentation](https://pytorch.org/docs/stable/_modules/torch/nn/init.html) and explore different initializers.\n",
    "1. Try accessing the model parameters after `net.apply(initialization)` and before `net(x)` to observe the shape of the model parameters. What changes? Why?\n",
    "1. Construct a multilayer perceptron containing a shared parameter layer and train it. During the training process, observe the model parameters and gradients of each layer.\n",
    "1. Why is sharing parameters a good idea?\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
